import json

from volcenginesdkarkruntime import Ark
from restful import DeepSeekResponse ,deepseek2openai



client = Ark(
    sk="AKLTOThkMDRiZWI2ZTNlNGY1YmE0NzRmNDE5OTQyYmU3ZWU",
    ak="WTJOaU5UWTFOV1V5TkRsaE5ETTNaR0ZsTURFNE9EazFNVGxsTm1JeVlqVQ==",
    api_key='78018cba-25ba-4e0b-ac6d-5ab1a2abecf6'
)

def doubao_chat(content) -> str:
    completion = client.chat.completions.create(
        model="ep-20240715005447-qdmcc",
        messages=[
            {"role": "system", "content": "你是豆包，是由字节跳动开发的 AI 人工智能助手"},
            {"role": "user", "content": content},
        ],
    )
    return completion.choices[0].message.content


def doubao_chat_stream(prompt):
    _stream = client.chat.completions.create(
        model="ep-20240715005447-qdmcc",
        messages=[
            {"role": "system", "content": "你是豆包，是由字节跳动开发的 AI 人工智能助手"},
            {"role": "user", "content": prompt},
        ],
        stream=True
    )
    for idx, item in enumerate(_stream):
        if item.choices[0].finish_reason=="stop":
            done = True
        else:
            done = False
        dsr = DeepSeekResponse(id=idx + 1, model=item.model, created_at=item.id,
                               response=item.choices[0].delta.content, done=done)
        new_res = deepseek2openai(dsr)
        yield new_res


    # for chunk in _stream:
    #     if not chunk.choices:
    #         continue
    #     # yield chunk.choices[0].delta.content
    #     yield chunk



if __name__ == '__main__':
    # for i in range(1):
    #     print(f"no_stream:{doubao_chat('关于“新质生产力”主题，生成具有三级标题的文章框架，在一级标题前增加“+”，在二级标题增加“++”，在三级标题前增加“+++”，检查是否正确的使用了标题级别表示符号“+”，只输出标题信息。')}")
    prompt = "中国的首都在哪里？"
    for i in doubao_chat_stream(prompt=prompt):
        print(type(i))
        # from types import SimpleNamespace
        # # 将 JSON 字符串转换为对象
        # i = json.loads(j, object_hook=lambda d: SimpleNamespace(**d))
        # print(type(i),i)
        # if i["choices"][0]["finish_reason"] == "stop":
        #     i["choices"][0]["delta"]["content"] = "stop"
        # print(i["choices"][0]["delta"]["content"])
        print(i.choices[0].delta.content)
        print(i.choices[0].finish_reason)
        # print(i.id)
        # print(i.created)
        # print(i.object)
    # res = doubao_chat_stream()f
    # stream = ''
    # for i in res:
    #     stream = stream + i
    # print(f"stream:{stream}")